ResultSetColumnNameHelperService.java

package com.opencsv;

/*
 * Copyright 2015 Scott Conway
 * 
 * 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.
 */

import org.apache.commons.lang3.ArrayUtils;
import org.apache.commons.lang3.ObjectUtils;
import org.apache.commons.lang3.StringUtils;

import java.io.IOException;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.*;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * Helper class for processing JDBC ResultSet objects allowing the user to
 * process a subset of columns and set custom header names.
 */
public class ResultSetColumnNameHelperService extends ResultSetHelperService implements ResultSetHelper {
    private String[] columnNames;
    private String[] columnHeaders;
    private final Map<String, Integer> columnNamePositionMap = new HashMap<>();
    private Locale errorLocale = Locale.getDefault();

    /**
     * Nullary constructor.
     */
    public ResultSetColumnNameHelperService() {
    }
    
    /**
     * Sets the locale for error messages.
     * @param errorLocale Locale for error messages. If null, the default locale
     *   is used.
     * @since 4.0
     */
    public void setErrorLocale(Locale errorLocale) {
        this.errorLocale = ObjectUtils.defaultIfNull(errorLocale, Locale.getDefault());
    }
    
    /**
     * Set the JDBC column names to use, and the header text for the CSV file
     * @param columnNames The JDBC column names to export, in the desired order
     * @param columnHeaders The column headers of the CSV file, in the desired order
     * @throws UnsupportedOperationException If the number of headers is different
     * than the number of columns, or if any of the columns or headers is blank
     * or null.
     */
    public void setColumnNames(String[] columnNames, String[] columnHeaders) {
        if (columnHeaders.length != columnNames.length) {
            throw new UnsupportedOperationException(ResourceBundle.getBundle(ICSVParser.DEFAULT_BUNDLE_NAME, errorLocale).getString("column.count.mismatch"));
        }
        if (hasInvalidValue(columnNames)) {
            throw new UnsupportedOperationException(ResourceBundle.getBundle(ICSVParser.DEFAULT_BUNDLE_NAME, errorLocale).getString("column.name.bogus"));
        }
        if (hasInvalidValue(columnHeaders)) {
            throw new UnsupportedOperationException(ResourceBundle.getBundle(ICSVParser.DEFAULT_BUNDLE_NAME, errorLocale).getString("header.name.bogus"));
        }
        this.columnNames = Arrays.copyOf(columnNames, columnNames.length);
        this.columnHeaders = Arrays.copyOf(columnHeaders, columnHeaders.length);
    }

    private boolean hasInvalidValue(String[] strings) {
        return Stream.of(strings).anyMatch(s -> StringUtils.isBlank(s));
    }

    /**
     * Returns the column names from the result set.
     * @param rs ResultSet
     * @return A string array containing the column names.
     * @throws SQLException Thrown by the result set.
     */
    @Override
    public String[] getColumnNames(ResultSet rs) throws SQLException {
        if (columnNamePositionMap.isEmpty()) {
            populateColumnData(rs);
        }
        return Arrays.copyOf(columnHeaders, columnHeaders.length);
    }

    private void populateColumnData(ResultSet rs) throws SQLException {
        String[] realColumnNames = super.getColumnNames(rs);

        if (columnNames == null) {
            columnNames = Arrays.copyOf(realColumnNames, realColumnNames.length);
            columnHeaders = Arrays.copyOf(realColumnNames, realColumnNames.length);
        }

        for (String name : columnNames) {
            int position = ArrayUtils.indexOf(realColumnNames, name);
            if (position == ArrayUtils.INDEX_NOT_FOUND) {
                throw new UnsupportedOperationException(String.format(ResourceBundle.getBundle(ICSVParser.DEFAULT_BUNDLE_NAME, errorLocale).getString("column.nonexistant"), name));
            }
            columnNamePositionMap.put(name, position);
        }
    }

    /**
     * Get all the column values from the result set.
     * @param rs The ResultSet containing the values.
     * @return String array containing all the column values.
     * @throws SQLException Thrown by the result set.
     * @throws IOException Thrown by the result set.
     */
    @Override
    public String[] getColumnValues(ResultSet rs) throws SQLException, IOException {
        if (columnNamePositionMap.isEmpty()) {
            populateColumnData(rs);
        }
        String[] realColumnValues = super.getColumnValues(rs, false, dateFormat, dateTimeFormat);
        return getColumnValueSubset(realColumnValues);
    }

    /**
     * Get all the column values from the result set.
     * @param rs The ResultSet containing the values.
     * @param trim Values should have white spaces trimmed.
     * @return String array containing all the column values.
     * @throws SQLException Thrown by the result set.
     * @throws IOException Thrown by the result set.
     */
    @Override
    public String[] getColumnValues(ResultSet rs, boolean trim) throws SQLException, IOException {
        if (columnNamePositionMap.isEmpty()) {
            populateColumnData(rs);
        }
        String[] realColumnValues = super.getColumnValues(rs, trim, dateFormat, dateTimeFormat);
        return getColumnValueSubset(realColumnValues);
    }

    /**
     * Get all the column values from the result set.
     * @param rs The ResultSet containing the values.
     * @param trim Values should have white spaces trimmed.
     * @param dateFormatString Format string for dates.
     * @param timeFormatString Format string for timestamps.
     * @return String array containing all the column values.
     * @throws SQLException Thrown by the result set.
     * @throws IOException Thrown by the result set.
     */
    @Override
    public String[] getColumnValues(ResultSet rs, boolean trim, String dateFormatString, String timeFormatString) throws SQLException, IOException {
        if (columnNamePositionMap.isEmpty()) {
            populateColumnData(rs);
        }
        String[] realColumnValues = super.getColumnValues(rs, trim, dateFormatString, timeFormatString);
        return getColumnValueSubset(realColumnValues);
    }

    private String[] getColumnValueSubset(String[] realColumnValues) {
        return Stream.of(columnNames)
                .map(c -> realColumnValues[columnNamePositionMap.get(c)])
                .collect(Collectors.toList())
                .toArray(ArrayUtils.EMPTY_STRING_ARRAY);
    }
}