com.baifendian.swordfish.execserver.engine.hive.HiveSqlExec.java Source code

Java tutorial

Introduction

Here is the source code for com.baifendian.swordfish.execserver.engine.hive.HiveSqlExec.java

Source

/*
 * Copyright (C) 2017 Baifendian Corporation
 *
 * 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 com.baifendian.swordfish.execserver.engine.hive;

import com.baifendian.swordfish.common.hive.service2.HiveService2Client;
import com.baifendian.swordfish.common.hive.service2.HiveService2ConnectionInfo;
import com.baifendian.swordfish.dao.DaoFactory;
import com.baifendian.swordfish.dao.enums.FlowStatus;
import com.baifendian.swordfish.dao.exception.DaoSemanticException;
import com.baifendian.swordfish.execserver.common.ExecResult;
import com.baifendian.swordfish.execserver.common.ResultCallback;
import com.baifendian.swordfish.execserver.utils.Constants;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.sql.SQLTimeoutException;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Date;
import java.util.List;
import java.util.function.Consumer;
import org.apache.commons.lang.StringUtils;
import org.apache.hive.jdbc.HiveConnection;
import org.apache.hive.jdbc.HiveStatement;
import org.apache.hive.service.cli.HiveSQLException;
import org.slf4j.Logger;

/**
 * Hive sql <p>
 */
public class HiveSqlExec {

    /**
     * ? 1000
     */
    private static int defaultQueryLimit = 1000;

    /**
     * {@link HiveUtil}
     */
    private final HiveUtil hiveUtil;

    /**
     * ?
     */
    private Consumer<List<String>> logHandler;

    /**
     * 
     */
    private String userName;

    /**
     * 
     */
    private Logger logger;

    /**
     * queue name
     */
    private String queueSQL;

    public HiveSqlExec(Consumer<List<String>> logHandler, String userName, Logger logger, String queue) {
        this.logHandler = logHandler;
        this.userName = userName;
        this.logger = logger;
        if (StringUtils.isEmpty(queue)) {
            queueSQL = null;
        } else {
            this.queueSQL = "SET mapreduce.job.queuename=" + queue;
        }
        this.hiveUtil = DaoFactory.getDaoInstance(HiveUtil.class);
    }

    /**
     *  sql ? ?, ?,  execute, ?
     *
     * @param createFuncs ?
     * @param sqls  sql
     * @param isContinue ?, ???
     * @param resultCallback , ?
     * @param queryLimit ?
     * @param remainTime ?, 
     */
    public boolean execute(List<String> createFuncs, List<String> sqls, boolean isContinue,
            ResultCallback resultCallback, Integer queryLimit, int remainTime) {

        // ?
        if (remainTime <= 0) {
            return false;
        }

        // ?
        queryLimit = (queryLimit != null) ? queryLimit : defaultQueryLimit;

        HiveConnection hiveConnection = null;
        Statement sta = null;
        Thread logThread = null;

        //  hive ?
        HiveService2ConnectionInfo hiveService2ConnectionInfo = hiveUtil.getHiveService2ConnectionInfo(userName);

        logger.info("execution connection information:{}", hiveService2ConnectionInfo);

        HiveService2Client hiveService2Client = hiveUtil.getHiveService2Client();

        try {
            try {
                hiveConnection = hiveService2Client.borrowClient(hiveService2ConnectionInfo);

                sta = hiveConnection.createStatement();
                //        sta.setQueryTimeout(remainTime);

                // 
                logThread = new Thread(new JdbcLogRunnable(sta));
                logThread.setDaemon(true);
                logThread.start();

                // set queue
                if (queueSQL != null) {
                    logger.info("hive queue : {}", queueSQL);
                    sta.execute(queueSQL);
                }

                //  function
                if (createFuncs != null) {
                    for (String createFunc : createFuncs) {
                        logger.info("hive create function sql: {}", createFunc);
                        sta.execute(createFunc);
                    }
                }
            } catch (Exception e) {
                logger.error("execute query exception", e);

                // , , ?
                handlerResults(0, sqls, FlowStatus.FAILED, resultCallback);

                return false;
            }

            //  sql ?
            for (int index = 0; index < sqls.size(); ++index) {
                String sql = sqls.get(index);

                Date startTime = new Date();

                logger.info("hive execute sql: {}", sql);

                ExecResult execResult = new ExecResult();
                execResult.setIndex(index);
                execResult.setStm(sql);

                try {
                    // ? query  show ?
                    if (HiveUtil.isTokQuery(sql) || HiveUtil.isLikeShowStm(sql)) {
                        sta.setMaxRows(queryLimit);
                        ResultSet res = sta.executeQuery(sql);

                        ResultSetMetaData resultSetMetaData = res.getMetaData();
                        int count = resultSetMetaData.getColumnCount();

                        List<String> colums = new ArrayList<>();
                        for (int i = 1; i <= count; i++) {
                            colums.add(resultSetMetaData.getColumnLabel(
                                    i)/*parseColumnName(resultSetMetaData.getColumnLabel(i), colums)*/);
                        }

                        execResult.setTitles(colums);

                        List<List<String>> datas = new ArrayList<>();

                        //  1,  query ?
                        if (count > 1 || HiveUtil.isTokQuery(sql)) {
                            while (res.next()) {
                                List<String> values = new ArrayList<>();
                                for (int i = 1; i <= count; ++i) {
                                    values.add(res.getString(i));
                                }

                                datas.add(values);
                            }
                        } else {
                            StringBuffer buffer = new StringBuffer();

                            while (res.next()) {
                                buffer.append(res.getString(1));
                                buffer.append("\n");
                            }

                            List<String> values = new ArrayList<>();
                            values.add(buffer.toString().trim());

                            datas.add(values);
                        }

                        execResult.setValues(datas);
                    } else {
                        sta.execute(sql);
                    }

                    // ??
                    execResult.setStatus(FlowStatus.SUCCESS);

                    // ?
                    if (resultCallback != null) {
                        Date endTime = new Date();
                        resultCallback.handleResult(execResult, startTime, endTime);
                    }
                } catch (SQLTimeoutException e) {
                    // sql 
                    logger.error("executeQuery timeout exception", e);

                    handlerResults(index, sqls, FlowStatus.FAILED, resultCallback);
                    return false;
                } catch (DaoSemanticException | HiveSQLException e) {
                    // 
                    logger.error("executeQuery exception", e);

                    if (isContinue) {
                        handlerResult(index, sql, FlowStatus.FAILED, resultCallback);
                    } else {
                        handlerResults(index, sqls, FlowStatus.FAILED, resultCallback);
                        return false;
                    }
                } catch (Exception e) {
                    // TTransport 
                    if (e.toString().contains("TTransportException")) {
                        logger.error("Get TTransportException return a client", e);
                        // ???
                        //            hiveService2Client.invalidateObject(hiveService2ConnectionInfo, hiveConnection);
                        handlerResults(index, sqls, FlowStatus.FAILED, resultCallback);
                        return false;
                    }

                    // socket 
                    if (e.toString().contains("SocketException")) {
                        logger.error("SocketException clear pool", e);
                        hiveService2Client.clear();
                        handlerResults(index, sqls, FlowStatus.FAILED, resultCallback);
                        return false;
                    }

                    logger.error("executeQuery exception", e);

                    if (isContinue) {
                        handlerResult(index, sql, FlowStatus.FAILED, resultCallback);
                    } else {
                        handlerResults(index, sqls, FlowStatus.FAILED, resultCallback);
                        return false;
                    }
                }
            }
        } finally {
            // 
            try {
                if (sta != null) {
                    sta.close();
                }
            } catch (Exception e) {
                logger.error("Catch an exception", e);
            }

            try {
                // 
                if (hiveConnection != null) {
                    // 
                    hiveConnection.close();

                    // , ??
                    hiveService2Client.returnClient(hiveService2ConnectionInfo, hiveConnection);
                }
            } catch (Exception e) {
                logger.error("Catch an exception", e);
            }

            // 
            try {
                if (logThread != null) {
                    logThread.interrupt();
                    logThread.join(HiveUtil.DEFAULT_QUERY_PROGRESS_THREAD_TIMEOUT);
                }
            } catch (Exception e) {
                //        logger.error("Catch an exception", e);
            }
        }

        return true;
    }

    /**
     * ?,  fromIndex 
     */
    private void handlerResults(int fromIndex, List<String> sqls, FlowStatus status,
            ResultCallback resultCallback) {
        for (int i = fromIndex; i < sqls.size(); ++i) {
            String sql = sqls.get(i);

            handlerResult(i, sql, status, resultCallback);
        }
    }

    /**
     * ???
     */
    private void handlerResult(int index, String sql, FlowStatus status, ResultCallback resultCallback) {
        Date now = new Date();

        ExecResult execResult = new ExecResult();

        execResult.setIndex(index);
        execResult.setStm(sql);
        execResult.setStatus(status);

        if (resultCallback != null) {
            // ?
            resultCallback.handleResult(execResult, now, now);
        }
    }

    /**
     * ? jdbc 
     */
    private class JdbcLogRunnable implements Runnable {

        private static final int DEFAULT_QUERY_PROGRESS_INTERVAL = 1000;

        private HiveStatement hiveStatement;
        private List<String> logs;

        public JdbcLogRunnable(Statement statement) {
            if (statement instanceof HiveStatement) {
                this.hiveStatement = (HiveStatement) statement;
            }

            this.logs = Collections.synchronizedList(new ArrayList<>());
        }

        @Override
        public void run() {
            if (hiveStatement == null) {
                return;
            }

            // ? flush 
            long preFlushTime = System.currentTimeMillis();

            synchronized (this) { // ??
                while (true) {
                    try {
                        for (String log : hiveStatement.getQueryLog()) {
                            logs.add(log);

                            long now = System.currentTimeMillis();

                            // ??
                            if (logs.size() >= Constants.defaultLogBufferSize
                                    || now - preFlushTime > Constants.defaultLogFlushInterval) {
                                preFlushTime = now;

                                // ?
                                logHandler.accept(logs);

                                logs.clear();
                            }
                        }

                        Thread.sleep(DEFAULT_QUERY_PROGRESS_INTERVAL);
                    } catch (SQLException e) {
                        return;
                    } catch (InterruptedException e) {
                        //            logger.error(e.getMessage(), e);
                        return;
                    } catch (Exception e) {
                        logger.error(e.getMessage(), e);
                        return;
                    } finally {
                        // ?
                        showRemainingLogsIfAny();
                    }
                }
            }
        }

        private void showRemainingLogsIfAny() {
            List<String> logsTemp;
            do {
                try {
                    logsTemp = hiveStatement.getQueryLog();
                } catch (Exception e) {
                    /*logger.error(e.getMessage(), e);*/
                    return;
                }

                for (String log : logsTemp) {
                    logs.add(log);

                    // ?
                    logHandler.accept(logs);

                    logs.clear();
                }
            } while (logsTemp.size() > 0);
        }
    }
}