com.huaweicloud.dws.client.handler.ScanActionHandler Maven / Gradle / Ivy

Go to download

package com.huaweicloud.dws.client.handler;

import com.huaweicloud.dws.client.DwsConfig;
import com.huaweicloud.dws.client.action.ScanAction;
import com.huaweicloud.dws.client.exception.DwsClientException;
import com.huaweicloud.dws.client.model.Record;
import com.huaweicloud.dws.client.model.TableName;
import com.huaweicloud.dws.client.model.TableSchema;
import com.huaweicloud.dws.client.op.Scan;
import com.huaweicloud.dws.client.util.JdbcUtil;
import com.huaweicloud.dws.client.util.LogUtil;
import com.huaweicloud.dws.client.worker.ConnectionProvider;
import lombok.extern.slf4j.Slf4j;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.BitSet;
import java.util.List;
import java.util.PrimitiveIterator;

/**
 * @ProjectName: dws-client
 * @Description:
 * @Date: 2023/10/19
 */
@Slf4j
public class ScanActionHandler extends AbstractActionHandler {
    private final DwsConfig config;

    public ScanActionHandler(ConnectionProvider connectionProvider, DwsConfig config) {
        super(connectionProvider);
        this.config = config;
    }

    @Override
    public void handle(ScanAction action) {
        Scan scan = action.getScan();
        if (scan == null) {
            log.debug("current action scan is empty.");
            action.getFuture().complete(null);
            return;
        }
        doHandleScanAction(action);
    }

    private void doHandleScanAction(ScanAction action) {
        try {
            Scan scan = action.getScan();
            Record record = scan.getRecord();
            TableSchema schema = record.getTableSchema();

            Connection connection = getConnection();
            List records = execute(schema, record.getTableSchema().getTableName(), scan, connection);

            if (!action.getFuture().isDone()) {
                action.getFuture().complete(records);
            }
        } catch (Exception e) {
            log.error("select executor execute fail. err = {}", e.getMessage());
            if (!action.getFuture().isDone()) {
                action.getFuture().completeExceptionally(e);
            }
        }
    }

    private List execute(TableSchema schema, TableName tableName, Scan scan, Connection connection)
        throws SQLException, DwsClientException {
        Record record = scan.getRecord();
        String sql = JdbcUtil.getScanFromStatement(schema, tableName, scan);
        LogUtil.withLogSwitch(config, () -> log.info("Scan sql:{}", sql));
        Record resultRecord = new Record(schema);
        BitSet columnMask = new BitSet(schema.getColumns().size());
        if (scan.getSelectedColumns() != null) {
            columnMask.or(scan.getSelectedColumns());
        } else {
            columnMask.set(0, schema.getColumns().size());
        }
        ResultSet rs = null;

        try (PreparedStatement ps =
            connection.prepareStatement(sql, ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY)) {
            int paramIndex = 0;
            for (PrimitiveIterator.OfInt it = record.getColumnBit().stream().iterator(); it.hasNext(); ) {
                Integer next = it.next();
                int type = record.getTableSchema().getColumn(next).getType();
                ps.setObject(++paramIndex, record.getValue(next), type);
            }
            ps.setFetchSize(scan.getFetchSize() > 0 ? scan.getFetchSize() : config.getScanFetchSize());
            ps.setQueryTimeout(scan.getTimeout() > 0 ? scan.getTimeout() : config.getScanTimeoutSeconds());
            rs = ps.executeQuery();
            return parseRecords(rs, resultRecord, columnMask, schema);
        } catch (Exception e) {
            log.error("Get executor execute fail. sql = {}, err = {}", sql, e.getMessage());
            throw DwsClientException.fromException(e);
        } finally {
            if (rs != null) {
                rs.close();
            }
        }
    }

    private List parseRecords(ResultSet rs, Record record, BitSet columnMask, TableSchema schema)
        throws SQLException {
        ArrayList records = new ArrayList<>();
        int count = rs.getFetchSize();
        while (rs.next() && count > 0) {
            Record copy = record.copy();
            int index = 0;
            for (PrimitiveIterator.OfInt it = columnMask.stream().iterator(); it.hasNext();) {
                int recordColumnIndex = it.next();
                // 将查询结果填充record
                fillRecord(copy, recordColumnIndex, rs, ++index, schema.getColumn(recordColumnIndex));
            }
            records.add(copy);
            count--;
        }
        return records;
    }
}