gobblin.source.extractor.extract.jdbc.JdbcExtractor.java Source code

Java tutorial

Introduction

Here is the source code for gobblin.source.extractor.extract.jdbc.JdbcExtractor.java

Source

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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 gobblin.source.extractor.extract.jdbc;

import java.io.IOException;
import java.sql.Blob;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.sql.Statement;
import java.sql.Types;
import java.text.ParseException;
import java.text.SimpleDateFormat;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Date;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

import org.apache.commons.codec.binary.Base64;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.base.Joiner;
import com.google.gson.Gson;
import com.google.gson.JsonArray;
import com.google.gson.JsonElement;
import com.google.gson.JsonObject;

import gobblin.configuration.ConfigurationKeys;
import gobblin.configuration.WorkUnitState;
import gobblin.password.PasswordManager;
import gobblin.source.extractor.DataRecordException;
import gobblin.source.extractor.exception.HighWatermarkException;
import gobblin.source.extractor.exception.RecordCountException;
import gobblin.source.extractor.exception.SchemaException;
import gobblin.source.extractor.extract.Command;
import gobblin.source.extractor.extract.CommandOutput;
import gobblin.source.extractor.extract.QueryBasedExtractor;
import gobblin.source.extractor.extract.SourceSpecificLayer;
import gobblin.source.extractor.extract.jdbc.JdbcCommand.JdbcCommandType;
import gobblin.source.extractor.resultset.RecordSetList;
import gobblin.source.extractor.schema.ColumnAttributes;
import gobblin.source.extractor.schema.ColumnNameCase;
import gobblin.source.extractor.schema.Schema;
import gobblin.source.extractor.utils.Utils;
import gobblin.source.extractor.watermark.Predicate;
import gobblin.source.extractor.watermark.WatermarkType;
import gobblin.source.workunit.WorkUnit;

/**
 * Extract data using JDBC protocol
 *
 * @author nveeramr
 */
public abstract class JdbcExtractor extends QueryBasedExtractor<JsonArray, JsonElement>
        implements SourceSpecificLayer<JsonArray, JsonElement>, JdbcSpecificLayer {
    private static final Gson gson = new Gson();
    private List<String> headerRecord;
    private boolean firstPull = true;
    private CommandOutput<?, ?> dataResponse = null;
    protected String extractSql;
    protected long sampleRecordCount;
    protected JdbcProvider jdbcSource;
    protected Connection dataConnection;
    protected int timeOut;
    private List<ColumnAttributes> columnAliasMap = new ArrayList<>();
    private Map<String, Schema> metadataColumnMap = new HashMap<>();
    private List<String> metadataColumnList = new ArrayList<>();
    private String inputColumnProjection;
    private String outputColumnProjection;
    private long totalRecordCount = 0;
    private boolean nextRecord = true;
    private int unknownColumnCounter = 1;

    private Logger log = LoggerFactory.getLogger(JdbcExtractor.class);

    /**
     * Metadata column mapping to lookup columns specified in input query
     *
     * @return metadata(schema) column mapping
     */
    public Map<String, Schema> getMetadataColumnMap() {
        return this.metadataColumnMap;
    }

    /**
     * @param metadataColumnMap metadata column mapping
     */
    public void setMetadataColumnMap(Map<String, Schema> metadataColumnMap) {
        this.metadataColumnMap = metadataColumnMap;
    }

    /**
     * Metadata column list
     *
     * @return metadata(schema) column list
     */
    public List<String> getMetadataColumnList() {
        return this.metadataColumnList;
    }

    /**
     * @param metadataColumnList metadata column list
     */
    public void setMetadataColumnList(List<String> metadataColumnList) {
        this.metadataColumnList = metadataColumnList;
    }

    /**
     * Sample Records specified in input query
     *
     * @return sample record count
     */
    public long getSampleRecordCount() {
        return this.sampleRecordCount;
    }

    /**
     * @param sampleRecordCount sample record count
     */
    public void setSampleRecordCount(long sampleRecordCount) {
        this.sampleRecordCount = sampleRecordCount;
    }

    /**
     * query to extract data from data source
     *
     * @return query
     */
    public String getExtractSql() {
        return this.extractSql;
    }

    /**
     * @param extractSql extract query
     */
    public void setExtractSql(String extractSql) {
        this.extractSql = extractSql;
    }

    /**
     * output column projection with aliases specified in input sql
     *
     * @return column projection
     */
    public String getOutputColumnProjection() {
        return this.outputColumnProjection;
    }

    /**
     * @param outputColumnProjection output column projection
     */
    public void setOutputColumnProjection(String outputColumnProjection) {
        this.outputColumnProjection = outputColumnProjection;
    }

    /**
     * input column projection with source columns specified in input sql
     *
     * @return column projection
     */
    public String getInputColumnProjection() {
        return this.inputColumnProjection;
    }

    /**
     * @param inputColumnProjection input column projection
     */
    public void setInputColumnProjection(String inputColumnProjection) {
        this.inputColumnProjection = inputColumnProjection;
    }

    /**
     * source column and alias mapping
     *
     * @return map of column name and alias name
     */
    public List<ColumnAttributes> getColumnAliasMap() {
        return this.columnAliasMap;
    }

    /**
     * add column and alias mapping
     *
     * @param columnAliasMap column alias mapping
     */
    public void addToColumnAliasMap(ColumnAttributes columnAliasMap) {
        this.columnAliasMap.add(columnAliasMap);
    }

    /**
     * check whether is first pull or not
     *
     * @return true, for the first run and it will be set to false after the
     *         first run
     */
    public boolean isFirstPull() {
        return this.firstPull;
    }

    /**
     * @param firstPull
     */
    public void setFirstPull(boolean firstPull) {
        this.firstPull = firstPull;
    }

    /**
     * Header record to convert csv to json
     *
     * @return header record with list of columns
     */
    protected List<String> getHeaderRecord() {
        return this.headerRecord;
    }

    /**
     * @param headerRecord list of column names
     */
    protected void setHeaderRecord(List<String> headerRecord) {
        this.headerRecord = headerRecord;
    }

    /**
     * @return connection timeout
     */
    public int getTimeOut() {
        return this.timeOut;
    }

    /**
     * @return true, if records available. Otherwise, false
     */
    public boolean hasNextRecord() {
        return this.nextRecord;
    }

    /**
     * @param nextRecord next Record
     */
    public void setNextRecord(boolean nextRecord) {
        this.nextRecord = nextRecord;
    }

    /**
     * @param timeOut connection timeout
     */
    @Override
    public void setTimeOut(int timeOut) {
        this.timeOut = timeOut;
    }

    /**
     * @return private static final Gson factory
     */
    public Gson getGson() {
        return this.gson;
    }

    public JdbcExtractor(WorkUnitState workUnitState) {
        super(workUnitState);
    }

    @Override
    public void extractMetadata(String schema, String entity, WorkUnit workUnit)
            throws SchemaException, IOException {
        this.log.info("Extract metadata using JDBC");
        String inputQuery = workUnitState.getProp(ConfigurationKeys.SOURCE_QUERYBASED_QUERY);
        String watermarkColumn = workUnitState.getProp(ConfigurationKeys.EXTRACT_DELTA_FIELDS_KEY);
        JsonObject defaultWatermark = this.getDefaultWatermark();
        String derivedWatermarkColumnName = defaultWatermark.get("columnName").getAsString();
        this.setSampleRecordCount(this.exractSampleRecordCountFromQuery(inputQuery));
        inputQuery = this.removeSampleClauseFromQuery(inputQuery);
        JsonArray targetSchema = new JsonArray();
        List<String> headerColumns = new ArrayList<>();

        try {
            List<Command> cmds = this.getSchemaMetadata(schema, entity);

            CommandOutput<?, ?> response = this.executePreparedSql(cmds);

            JsonArray array = this.getSchema(response);

            this.buildMetadataColumnMap(array);
            this.parseInputQuery(inputQuery);
            List<String> sourceColumns = this.getMetadataColumnList();

            for (ColumnAttributes colMap : this.columnAliasMap) {
                String alias = colMap.getAliasName();
                String columnName = colMap.getColumnName();
                String sourceColumnName = colMap.getSourceColumnName();
                if (this.isMetadataColumn(columnName, sourceColumns)) {
                    String targetColumnName = this.getTargetColumnName(columnName, alias);
                    Schema obj = this.getUpdatedSchemaObject(columnName, alias, targetColumnName);
                    String jsonStr = gson.toJson(obj);
                    JsonObject jsonObject = gson.fromJson(jsonStr, JsonObject.class).getAsJsonObject();
                    targetSchema.add(jsonObject);
                    headerColumns.add(targetColumnName);
                    this.columnList.add(sourceColumnName);
                }
            }

            if (this.hasMultipleWatermarkColumns(watermarkColumn)) {
                this.columnList.add(derivedWatermarkColumnName);
                headerColumns.add(derivedWatermarkColumnName);
                targetSchema.add(defaultWatermark);
                this.workUnitState.setProp(ConfigurationKeys.EXTRACT_DELTA_FIELDS_KEY, derivedWatermarkColumnName);
            }

            String outputColProjection = Joiner.on(",").useForNull("null").join(this.columnList);
            outputColProjection = outputColProjection.replace(derivedWatermarkColumnName,
                    Utils.getCoalesceColumnNames(watermarkColumn) + " AS " + derivedWatermarkColumnName);
            this.setOutputColumnProjection(outputColProjection);
            String extractQuery = this.getExtractQuery(schema, entity, inputQuery);

            this.setHeaderRecord(headerColumns);
            this.setOutputSchema(targetSchema);
            this.setExtractSql(extractQuery);
            // this.workUnit.getProp(ConfigurationKeys.EXTRACT_TABLE_NAME_KEY,
            // this.escapeCharsInColumnName(this.workUnit.getProp(ConfigurationKeys.SOURCE_ENTITY),
            // ConfigurationKeys.ESCAPE_CHARS_IN_COLUMN_NAME, "_"));
            this.log.info("Schema:" + targetSchema);
            this.log.info("Extract query: " + this.getExtractSql());
        } catch (RuntimeException | IOException | SchemaException e) {
            throw new SchemaException("Failed to get metadata using JDBC; error - " + e.getMessage(), e);
        }
    }

    /**
     * Build/Format input query in the required format
     *
     * @param schema
     * @param entity
     * @param inputQuery
     * @return formatted extract query
     */
    private String getExtractQuery(String schema, String entity, String inputQuery) {
        String inputColProjection = this.getInputColumnProjection();
        String outputColProjection = this.getOutputColumnProjection();
        String query = inputQuery;
        if (query == null) {
            // if input query is null, build the query from metadata
            query = "SELECT " + outputColProjection + " FROM " + schema + "." + entity;
        } else {
            // replace input column projection with output column projection
            if (StringUtils.isNotBlank(inputColProjection)) {
                query = query.replace(inputColProjection, outputColProjection);
            }
        }

        String watermarkPredicateSymbol = ConfigurationKeys.DEFAULT_SOURCE_QUERYBASED_WATERMARK_PREDICATE_SYMBOL;
        if (!query.contains(watermarkPredicateSymbol)) {
            query = SqlQueryUtils.addPredicate(query, watermarkPredicateSymbol);
        }
        return query;
    }

    /**
     * Update schema of source column Update column name with target column
     * name/alias Update watermark, nullable and primary key flags
     *
     * @param sourceColumnName
     * @param targetColumnName
     * @return schema object of a column
     */
    private Schema getUpdatedSchemaObject(String sourceColumnName, String alias, String targetColumnName) {
        // Check for source column and alias
        Schema obj = this.getMetadataColumnMap().get(sourceColumnName.toLowerCase());
        if (obj == null && alias != null) {
            obj = this.getMetadataColumnMap().get(alias.toLowerCase());
        }

        if (obj == null) {
            obj = getCustomColumnSchema(targetColumnName);
        } else {
            String watermarkColumn = this.workUnitState.getProp(ConfigurationKeys.EXTRACT_DELTA_FIELDS_KEY);
            String primarykeyColumn = this.workUnitState.getProp(ConfigurationKeys.EXTRACT_PRIMARY_KEY_FIELDS_KEY);
            boolean isMultiColumnWatermark = this.hasMultipleWatermarkColumns(watermarkColumn);

            obj.setColumnName(targetColumnName);
            boolean isWatermarkColumn = this.isWatermarkColumn(watermarkColumn, sourceColumnName);
            if (isWatermarkColumn) {
                this.updateDeltaFieldConfig(sourceColumnName, targetColumnName);
            } else if (alias != null) {
                // Check for alias
                isWatermarkColumn = this.isWatermarkColumn(watermarkColumn, alias);
                this.updateDeltaFieldConfig(alias, targetColumnName);
            }

            // If there is only one watermark column, then consider it as a
            // watermark. Otherwise add a default watermark column in the end
            if (!isMultiColumnWatermark) {
                obj.setWaterMark(isWatermarkColumn);
            }

            // override all columns to nullable except primary key and watermark
            // columns
            if ((isWatermarkColumn && !isMultiColumnWatermark)
                    || this.getPrimarykeyIndex(primarykeyColumn, sourceColumnName) > 0) {
                obj.setNullable(false);
            } else {
                obj.setNullable(true);
            }

            // set primary key index for all the primary key fields
            int primarykeyIndex = this.getPrimarykeyIndex(primarykeyColumn, sourceColumnName);
            if (primarykeyIndex > 0 && (!sourceColumnName.equalsIgnoreCase(targetColumnName))) {
                this.updatePrimaryKeyConfig(sourceColumnName, targetColumnName);
            }

            obj.setPrimaryKey(primarykeyIndex);
        }
        return obj;
    }

    /**
     * Get target column name if column is not found in metadata, then name it
     * as unknown column If alias is not found, target column is nothing but
     * source column
     *
     * @param sourceColumnName
     * @param alias
     * @return targetColumnName
     */
    private String getTargetColumnName(String sourceColumnName, String alias) {
        String targetColumnName = alias;
        Schema obj = this.getMetadataColumnMap().get(sourceColumnName.toLowerCase());
        if (obj == null) {
            targetColumnName = (targetColumnName == null ? "unknown" + this.unknownColumnCounter
                    : targetColumnName);
            this.unknownColumnCounter++;
        } else {
            targetColumnName = (StringUtils.isNotBlank(targetColumnName) ? targetColumnName : sourceColumnName);
        }
        targetColumnName = this.toCase(targetColumnName);
        return Utils.escapeSpecialCharacters(targetColumnName, ConfigurationKeys.ESCAPE_CHARS_IN_COLUMN_NAME, "_");
    }

    /**
     * Build metadata column map with column name and column schema object.
     * Build metadata column list with list columns in metadata
     *
     * @param array Schema of all columns
     */
    private void buildMetadataColumnMap(JsonArray array) {
        if (array != null) {
            for (JsonElement columnElement : array) {
                Schema schemaObj = gson.fromJson(columnElement, Schema.class);
                String columnName = schemaObj.getColumnName();
                this.metadataColumnMap.put(columnName.toLowerCase(), schemaObj);
                this.metadataColumnList.add(columnName.toLowerCase());
            }
        }
    }

    /**
     * Update water mark column property if there is an alias defined in query
     *
     * @param srcColumnName source column name
     * @param tgtColumnName target column name
     */
    private void updateDeltaFieldConfig(String srcColumnName, String tgtColumnName) {
        if (this.workUnitState.contains(ConfigurationKeys.EXTRACT_DELTA_FIELDS_KEY)) {
            String watermarkCol = this.workUnitState.getProp(ConfigurationKeys.EXTRACT_DELTA_FIELDS_KEY);
            this.workUnitState.setProp(ConfigurationKeys.EXTRACT_DELTA_FIELDS_KEY,
                    watermarkCol.replaceAll(srcColumnName, tgtColumnName));
        }
    }

    /**
     * Update primary key column property if there is an alias defined in query
     *
     * @param srcColumnName source column name
     * @param tgtColumnName target column name
     */
    private void updatePrimaryKeyConfig(String srcColumnName, String tgtColumnName) {
        if (this.workUnitState.contains(ConfigurationKeys.EXTRACT_PRIMARY_KEY_FIELDS_KEY)) {
            String primarykey = this.workUnitState.getProp(ConfigurationKeys.EXTRACT_PRIMARY_KEY_FIELDS_KEY);
            this.workUnitState.setProp(ConfigurationKeys.EXTRACT_PRIMARY_KEY_FIELDS_KEY,
                    primarykey.replaceAll(srcColumnName, tgtColumnName));
        }
    }

    /**
     * If input query is null or '*' in the select list, consider all columns.
     *
     * @return true, to select all colums. else, false.
     */
    private boolean isSelectAllColumns() {
        String columnProjection = this.getInputColumnProjection();
        if (columnProjection == null || columnProjection.trim().equals("*") || columnProjection.contains(".*")) {
            return true;
        }
        return false;
    }

    /**
     * Parse query provided in pull file Set input column projection - column
     * projection in the input query Set columnAlias map - column and its alias
     * mentioned in input query
     *
     * @param query input query
     */
    private void parseInputQuery(String query) {
        List<String> projectedColumns = new ArrayList<>();
        if (StringUtils.isNotBlank(query)) {
            String queryLowerCase = query.toLowerCase();
            int startIndex = queryLowerCase.indexOf("select ") + 7;
            int endIndex = queryLowerCase.indexOf(" from ");
            if (startIndex >= 0 && endIndex >= 0) {
                String columnProjection = query.substring(startIndex, endIndex);
                this.setInputColumnProjection(columnProjection);
                // parse the select list
                StringBuffer sb = new StringBuffer();
                int bracketCount = 0;
                for (int i = 0; i < columnProjection.length(); i++) {
                    char c = columnProjection.charAt(i);
                    if (c == '(') {
                        bracketCount++;
                    }

                    if (c == ')') {
                        bracketCount--;
                    }

                    if (bracketCount != 0) {
                        sb.append(c);
                    } else {
                        if (c != ',') {
                            sb.append(c);
                        } else {
                            projectedColumns.add(sb.toString());
                            sb = new StringBuffer();
                        }
                    }
                }
                projectedColumns.add(sb.toString());
            }
        }

        if (this.isSelectAllColumns()) {
            List<String> columnList = this.getMetadataColumnList();
            for (String columnName : columnList) {
                ColumnAttributes col = new ColumnAttributes();
                col.setColumnName(columnName);
                col.setAliasName(columnName);
                col.setSourceColumnName(columnName);
                this.addToColumnAliasMap(col);
            }
        } else {
            for (String projectedColumn : projectedColumns) {
                String column = projectedColumn.trim();
                String alias = null;
                String sourceColumn = column;
                int spaceOccurences = StringUtils.countMatches(column.trim(), " ");
                if (spaceOccurences > 0) {
                    // separate column and alias if they are separated by "as"
                    // or space
                    int lastSpaceIndex = column.toLowerCase().lastIndexOf(" as ");
                    sourceColumn = column.substring(0, lastSpaceIndex);
                    alias = column.substring(lastSpaceIndex + 4);
                }

                // extract column name if projection has table name in it
                String columnName = sourceColumn;
                if (sourceColumn.contains(".")) {
                    columnName = sourceColumn.substring(sourceColumn.indexOf(".") + 1);
                }

                ColumnAttributes col = new ColumnAttributes();
                col.setColumnName(columnName);
                col.setAliasName(alias);
                col.setSourceColumnName(sourceColumn);
                this.addToColumnAliasMap(col);
            }
        }
    }

    /**
     * Execute query using JDBC simple Statement Set fetch size
     *
     * @param cmds commands - query, fetch size
     * @return JDBC ResultSet
     * @throws Exception
     */
    private CommandOutput<?, ?> executeSql(List<Command> cmds) {
        String query = null;
        int fetchSize = 0;

        for (Command cmd : cmds) {
            if (cmd instanceof JdbcCommand) {
                JdbcCommandType type = (JdbcCommandType) cmd.getCommandType();
                switch (type) {
                case QUERY:
                    query = cmd.getParams().get(0);
                    break;
                case FETCHSIZE:
                    fetchSize = Integer.parseInt(cmd.getParams().get(0));
                    break;
                default:
                    this.log.error("Command " + type.toString() + " not recognized");
                    break;
                }
            }
        }

        this.log.info("Executing query:" + query);
        ResultSet resultSet = null;
        try {
            this.jdbcSource = createJdbcSource();
            this.dataConnection = this.jdbcSource.getConnection();
            Statement statement = this.dataConnection.createStatement();

            if (fetchSize != 0 && this.getExpectedRecordCount() > 2000) {
                statement.setFetchSize(fetchSize);
            }
            final boolean status = statement.execute(query);
            if (status == false) {
                this.log.error("Failed to execute sql:" + query);
            }
            resultSet = statement.getResultSet();
        } catch (Exception e) {
            this.log.error("Failed to execute sql:" + query + " ;error-" + e.getMessage(), e);
        }

        CommandOutput<JdbcCommand, ResultSet> output = new JdbcCommandOutput();
        output.put((JdbcCommand) cmds.get(0), resultSet);
        return output;
    }

    /**
     * Execute query using JDBC PreparedStatement to pass query parameters Set
     * fetch size
     *
     * @param cmds commands - query, fetch size, query parameters
     * @return JDBC ResultSet
     * @throws Exception
     */
    private CommandOutput<?, ?> executePreparedSql(List<Command> cmds) {
        String query = null;
        List<String> queryParameters = null;
        int fetchSize = 0;

        for (Command cmd : cmds) {
            if (cmd instanceof JdbcCommand) {
                JdbcCommandType type = (JdbcCommandType) cmd.getCommandType();
                switch (type) {
                case QUERY:
                    query = cmd.getParams().get(0);
                    break;
                case QUERYPARAMS:
                    queryParameters = cmd.getParams();
                    break;
                case FETCHSIZE:
                    fetchSize = Integer.parseInt(cmd.getParams().get(0));
                    break;
                default:
                    this.log.error("Command " + type.toString() + " not recognized");
                    break;
                }
            }
        }

        this.log.info("Executing query:" + query);
        ResultSet resultSet = null;
        try {
            this.jdbcSource = createJdbcSource();
            this.dataConnection = this.jdbcSource.getConnection();

            PreparedStatement statement = this.dataConnection.prepareStatement(query, ResultSet.TYPE_FORWARD_ONLY,
                    ResultSet.CONCUR_READ_ONLY);

            int parameterPosition = 1;
            if (queryParameters != null && queryParameters.size() > 0) {
                for (String parameter : queryParameters) {
                    statement.setString(parameterPosition, parameter);
                    parameterPosition++;
                }
            }
            if (fetchSize != 0) {
                statement.setFetchSize(fetchSize);
            }
            final boolean status = statement.execute();
            if (status == false) {
                this.log.error("Failed to execute sql:" + query);
            }
            resultSet = statement.getResultSet();

        } catch (Exception e) {
            this.log.error("Failed to execute sql:" + query + " ;error-" + e.getMessage(), e);
        }

        CommandOutput<JdbcCommand, ResultSet> output = new JdbcCommandOutput();
        output.put((JdbcCommand) cmds.get(0), resultSet);
        return output;
    }

    /**
     * Create JDBC source to get connection
     *
     * @return JDBCSource
     */
    protected JdbcProvider createJdbcSource() {
        String driver = this.workUnitState.getProp(ConfigurationKeys.SOURCE_CONN_DRIVER);
        String userName = this.workUnitState.getProp(ConfigurationKeys.SOURCE_CONN_USERNAME);
        String password = PasswordManager.getInstance(this.workUnitState)
                .readPassword(this.workUnitState.getProp(ConfigurationKeys.SOURCE_CONN_PASSWORD));
        String connectionUrl = this.getConnectionUrl();

        String proxyHost = this.workUnitState.getProp(ConfigurationKeys.SOURCE_CONN_USE_PROXY_URL);
        int proxyPort = this.workUnitState.getProp(ConfigurationKeys.SOURCE_CONN_USE_PROXY_PORT) != null
                ? this.workUnitState.getPropAsInt(ConfigurationKeys.SOURCE_CONN_USE_PROXY_PORT)
                : -1;

        if (this.jdbcSource == null || this.jdbcSource.isClosed()) {
            this.jdbcSource = new JdbcProvider(driver, connectionUrl, userName, password, 1, this.getTimeOut(),
                    "DEFAULT", proxyHost, proxyPort);
            return this.jdbcSource;
        } else {
            return this.jdbcSource;
        }
    }

    @Override
    public long getMaxWatermark(String schema, String entity, String watermarkColumn, List<Predicate> predicateList,
            String watermarkSourceFormat) throws HighWatermarkException {
        this.log.info("Get high watermark using JDBC");
        long calculatedHighWatermark = ConfigurationKeys.DEFAULT_WATERMARK_VALUE;

        try {
            List<Command> cmds = this.getHighWatermarkMetadata(schema, entity, watermarkColumn, predicateList);
            CommandOutput<?, ?> response = this.executeSql(cmds);
            calculatedHighWatermark = this.getHighWatermark(response, watermarkColumn, watermarkSourceFormat);
            return calculatedHighWatermark;
        } catch (Exception e) {
            throw new HighWatermarkException("Failed to get high watermark using JDBC; error - " + e.getMessage(),
                    e);
        }
    }

    @Override
    public long getSourceCount(String schema, String entity, WorkUnit workUnit, List<Predicate> predicateList)
            throws RecordCountException {
        this.log.info("Get source record count using JDBC");
        long count = 0;
        try {
            List<Command> cmds = this.getCountMetadata(schema, entity, workUnit, predicateList);
            CommandOutput<?, ?> response = this.executeSql(cmds);
            count = this.getCount(response);
            this.log.info("Source record count:" + count);
            return count;
        } catch (Exception e) {
            throw new RecordCountException(
                    "Failed to get source record count using JDBC; error - " + e.getMessage(), e);
        }
    }

    @Override
    public Iterator<JsonElement> getRecordSet(String schema, String entity, WorkUnit workUnit,
            List<Predicate> predicateList) throws DataRecordException, IOException {
        Iterator<JsonElement> rs = null;
        List<Command> cmds;
        try {
            if (isFirstPull()) {
                this.log.info("Get data recordset using JDBC");
                cmds = this.getDataMetadata(schema, entity, workUnit, predicateList);
                this.dataResponse = this.executePreparedSql(cmds);
                this.setFirstPull(false);
            }

            rs = this.getData(this.dataResponse);
            return rs;
        } catch (Exception e) {
            throw new DataRecordException("Failed to get record set using JDBC; error - " + e.getMessage(), e);
        }
    }

    @Override
    public JsonArray getSchema(CommandOutput<?, ?> response) throws SchemaException, IOException {
        this.log.debug("Extract schema from resultset");
        ResultSet resultset = null;
        Iterator<ResultSet> itr = (Iterator<ResultSet>) response.getResults().values().iterator();
        if (itr.hasNext()) {
            resultset = itr.next();
        } else {
            throw new SchemaException("Failed to get schema from database - Resultset has no records");
        }

        JsonArray fieldJsonArray = new JsonArray();
        try {
            while (resultset.next()) {
                Schema schema = new Schema();
                String columnName = resultset.getString(1);
                schema.setColumnName(columnName);

                String dataType = resultset.getString(2);
                String elementDataType = "string";
                List<String> mapSymbols = null;
                JsonObject newDataType = this.convertDataType(columnName, dataType, elementDataType, mapSymbols);

                schema.setDataType(newDataType);
                schema.setLength(resultset.getLong(3));
                schema.setPrecision(resultset.getInt(4));
                schema.setScale(resultset.getInt(5));
                schema.setNullable(resultset.getBoolean(6));
                schema.setFormat(resultset.getString(7));
                schema.setComment(resultset.getString(8));
                schema.setDefaultValue(null);
                schema.setUnique(false);

                String jsonStr = gson.toJson(schema);
                JsonObject obj = gson.fromJson(jsonStr, JsonObject.class).getAsJsonObject();
                fieldJsonArray.add(obj);
            }
        } catch (Exception e) {
            throw new SchemaException("Failed to get schema from database; error - " + e.getMessage(), e);
        }

        return fieldJsonArray;
    }

    @Override
    public long getHighWatermark(CommandOutput<?, ?> response, String watermarkColumn, String watermarkColumnFormat)
            throws HighWatermarkException {
        this.log.debug("Extract high watermark from resultset");
        ResultSet resultset = null;
        Iterator<ResultSet> itr = (Iterator<ResultSet>) response.getResults().values().iterator();
        if (itr.hasNext()) {
            resultset = itr.next();
        } else {
            throw new HighWatermarkException(
                    "Failed to get high watermark from database - Resultset has no records");
        }

        Long HighWatermark;
        try {
            String watermark;
            if (resultset.next()) {
                watermark = resultset.getString(1);
            } else {
                watermark = null;
            }

            if (watermark == null) {
                return ConfigurationKeys.DEFAULT_WATERMARK_VALUE;
            }

            if (watermarkColumnFormat != null) {
                SimpleDateFormat inFormat = new SimpleDateFormat(watermarkColumnFormat);
                Date date = null;
                try {
                    date = inFormat.parse(watermark);
                } catch (ParseException e) {
                    this.log.error("ParseException: " + e.getMessage(), e);
                }
                SimpleDateFormat outFormat = new SimpleDateFormat("yyyyMMddHHmmss");
                HighWatermark = Long.parseLong(outFormat.format(date));
            } else {
                HighWatermark = Long.parseLong(watermark);
            }
        } catch (Exception e) {
            throw new HighWatermarkException(
                    "Failed to get high watermark from database; error - " + e.getMessage(), e);
        }

        return HighWatermark;
    }

    @Override
    public long getCount(CommandOutput<?, ?> response) throws RecordCountException {
        this.log.debug("Extract source record count from resultset");
        ResultSet resultset = null;
        long count = 0;
        Iterator<ResultSet> itr = (Iterator<ResultSet>) response.getResults().values().iterator();
        if (itr.hasNext()) {
            resultset = itr.next();

            try {
                if (resultset.next()) {
                    count = resultset.getLong(1);
                }
            } catch (Exception e) {
                throw new RecordCountException(
                        "Failed to get source record count from database; error - " + e.getMessage(), e);
            }
        } else {
            throw new RuntimeException(
                    "Failed to get source record count from database - Resultset has no records");
        }

        return count;
    }

    @Override
    public Iterator<JsonElement> getData(CommandOutput<?, ?> response) throws DataRecordException, IOException {
        this.log.debug("Extract data records from resultset");

        RecordSetList<JsonElement> recordSet = this.getNewRecordSetList();

        if (response == null || !this.hasNextRecord()) {
            return recordSet.iterator();
        }

        ResultSet resultset = null;
        Iterator<ResultSet> itr = (Iterator<ResultSet>) response.getResults().values().iterator();
        if (itr.hasNext()) {
            resultset = itr.next();
        } else {
            throw new DataRecordException(
                    "Failed to get source record count from database - Resultset has no records");
        }

        try {
            final ResultSetMetaData resultsetMetadata = resultset.getMetaData();

            int batchSize = this.workUnitState.getPropAsInt(ConfigurationKeys.SOURCE_QUERYBASED_FETCH_SIZE, 0);
            batchSize = (batchSize == 0 ? ConfigurationKeys.DEFAULT_SOURCE_FETCH_SIZE : batchSize);

            int recordCount = 0;
            while (resultset.next()) {

                final int numColumns = resultsetMetadata.getColumnCount();
                JsonObject jsonObject = new JsonObject();

                for (int i = 1; i < numColumns + 1; i++) {
                    final String columnName = this.getHeaderRecord().get(i - 1);
                    jsonObject.addProperty(columnName, parseColumnAsString(resultset, resultsetMetadata, i));

                }

                recordSet.add(jsonObject);

                recordCount++;
                this.totalRecordCount++;

                // Insert records in record set until it reaches the batch size
                if (recordCount >= batchSize) {
                    this.log.info("Total number of records processed so far: " + this.totalRecordCount);
                    return recordSet.iterator();
                }
            }
            this.setNextRecord(false);
            this.log.info("Total number of records processed so far: " + this.totalRecordCount);
            return recordSet.iterator();
        } catch (Exception e) {
            throw new DataRecordException("Failed to get records from database; error - " + e.getMessage(), e);
        }
    }

    /*
     * For Blob data, need to get the bytes and use base64 encoding to encode the byte[]
     * When reading from the String, need to use base64 decoder
     *     String tmp = ... ( get the String value )
     *     byte[] foo = Base64.decodeBase64(tmp);
     */
    private String readBlobAsString(Blob logBlob) throws SQLException {
        if (logBlob == null) {
            return StringUtils.EMPTY;
        }

        byte[] ba = logBlob.getBytes(1L, (int) (logBlob.length()));

        if (ba == null) {
            return StringUtils.EMPTY;
        }
        String baString = Base64.encodeBase64String(ba);
        return baString;
    }

    /**
     * HACK: there is a bug in the MysqlExtractor where tinyint columns are always treated as ints.
     * There are MySQL jdbc driver setting (tinyInt1isBit=true and transformedBitIsBoolean=false) that
     * can cause tinyint(1) columns to be treated as BIT/BOOLEAN columns. The default behavior is to
     * treat tinyint(1) as BIT.
     *
     * Currently, {@link MysqlExtractor#getDataTypeMap()} uses the information_schema to check types.
     * That does not do the above conversion. {@link #parseColumnAsString(ResultSet, ResultSetMetaData, int)}
     * which does the above type mapping.
     *
     * On the other hand, SqlServerExtractor treats BIT columns as Booleans. So we can be in a bind
     * where sometimes BIT has to be converted to an int (for backwards compatibility in MySQL) and
     * sometimes to a Boolean (for SqlServer).
     *
     * This function adds configurable behavior depending on the Extractor type.
     **/
    protected boolean convertBitToBoolean() {
        return true;
    }

    private String parseColumnAsString(final ResultSet resultset, final ResultSetMetaData resultsetMetadata, int i)
            throws SQLException {

        if (isBlob(resultsetMetadata.getColumnType(i))) {
            return readBlobAsString(resultset.getBlob(i));
        }
        if ((resultsetMetadata.getColumnType(i) == Types.BIT || resultsetMetadata.getColumnType(i) == Types.BOOLEAN)
                && convertBitToBoolean()) {
            return Boolean.toString(resultset.getBoolean(i));
        }
        return resultset.getString(i);
    }

    private static boolean isBlob(int columnType) {
        return columnType == Types.LONGVARBINARY || columnType == Types.BINARY;
    }

    protected static Command getCommand(String query, JdbcCommandType commandType) {
        return new JdbcCommand().build(Arrays.asList(query), commandType);
    }

    protected static Command getCommand(int fetchSize, JdbcCommandType commandType) {
        return new JdbcCommand().build(Arrays.asList(Integer.toString(fetchSize)), commandType);
    }

    protected static Command getCommand(List<String> params, JdbcCommandType commandType) {
        return new JdbcCommand().build(params, commandType);
    }

    /**
     * Concatenate all predicates with "and" clause
     *
     * @param predicateList list of predicate(filter) conditions
     * @return predicate
     */
    protected String concatPredicates(List<Predicate> predicateList) {
        List<String> conditions = new ArrayList<>();
        for (Predicate predicate : predicateList) {
            conditions.add(predicate.getCondition());
        }
        return Joiner.on(" and ").skipNulls().join(conditions);
    }

    /**
     * Schema of default watermark column-required if there are multiple watermarks
     *
     * @return column schema
     */
    private JsonObject getDefaultWatermark() {
        Schema schema = new Schema();
        String dataType;
        String columnName = "derivedwatermarkcolumn";

        schema.setColumnName(columnName);

        WatermarkType wmType = WatermarkType.valueOf(this.workUnitState
                .getProp(ConfigurationKeys.SOURCE_QUERYBASED_WATERMARK_TYPE, "TIMESTAMP").toUpperCase());
        switch (wmType) {
        case TIMESTAMP:
            dataType = "timestamp";
            break;
        case DATE:
            dataType = "date";
            break;
        default:
            dataType = "int";
            break;
        }

        String elementDataType = "string";
        List<String> mapSymbols = null;
        JsonObject newDataType = this.convertDataType(columnName, dataType, elementDataType, mapSymbols);
        schema.setDataType(newDataType);
        schema.setWaterMark(true);
        schema.setPrimaryKey(0);
        schema.setLength(0);
        schema.setPrecision(0);
        schema.setScale(0);
        schema.setNullable(false);
        schema.setFormat(null);
        schema.setComment("Default watermark column");
        schema.setDefaultValue(null);
        schema.setUnique(false);

        String jsonStr = gson.toJson(schema);
        JsonObject obj = gson.fromJson(jsonStr, JsonObject.class).getAsJsonObject();
        return obj;
    }

    /**
     * Schema of a custom column - required if column not found in metadata
     *
     * @return column schema
     */
    private Schema getCustomColumnSchema(String columnName) {
        Schema schema = new Schema();
        String dataType = "string";
        schema.setColumnName(columnName);
        String elementDataType = "string";
        List<String> mapSymbols = null;
        JsonObject newDataType = this.convertDataType(columnName, dataType, elementDataType, mapSymbols);
        schema.setDataType(newDataType);
        schema.setWaterMark(false);
        schema.setPrimaryKey(0);
        schema.setLength(0);
        schema.setPrecision(0);
        schema.setScale(0);
        schema.setNullable(true);
        schema.setFormat(null);
        schema.setComment("Custom column");
        schema.setDefaultValue(null);
        schema.setUnique(false);
        return schema;
    }

    /**
     * New record set for iterator
     *
     * @return RecordSetList
     */
    private static RecordSetList<JsonElement> getNewRecordSetList() {
        return new RecordSetList<>();
    }

    /**
     * Change the column name case to upper, lower or nochange; Default nochange
     *
     * @return column name with the required case
     */
    private String toCase(String targetColumnName) {
        String columnName = targetColumnName;
        ColumnNameCase caseType = ColumnNameCase.valueOf(this.workUnitState
                .getProp(ConfigurationKeys.SOURCE_COLUMN_NAME_CASE, ConfigurationKeys.DEFAULT_COLUMN_NAME_CASE)
                .toUpperCase());
        switch (caseType) {
        case TOUPPER:
            columnName = targetColumnName.toUpperCase();
            break;
        case TOLOWER:
            columnName = targetColumnName.toLowerCase();
            break;
        default:
            columnName = targetColumnName;
            break;
        }
        return columnName;
    }

    @Override
    public void closeConnection() throws Exception {
        if (this.dataConnection != null) {
            try {
                this.dataConnection.close();
            } catch (SQLException e) {
                this.log.error("Failed to close connection ;error-" + e.getMessage(), e);
            }
        }

        this.jdbcSource.close();
    }
}