com.huaweicloud.dws.client.executor.GetExecutor Maven / Gradle / Ivy

Go to download

package com.huaweicloud.dws.client.executor;

import com.huaweicloud.dws.client.DwsConfig;
import com.huaweicloud.dws.client.exception.DwsClientException;
import com.huaweicloud.dws.client.model.Record;
import com.huaweicloud.dws.client.model.TableSchema;
import com.huaweicloud.dws.client.op.Get;
import com.huaweicloud.dws.client.util.JdbcUtil;
import com.huaweicloud.dws.client.util.LogUtil;
import lombok.extern.slf4j.Slf4j;

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

import static com.huaweicloud.dws.client.handler.AbstractActionHandler.fillRecord;

/**
 * @ProjectName: dws-client
 * @Description: get操作执行器
 * @Date: 2023/10/09
 */
@Slf4j
public class GetExecutor {
    public Map execute(TableSchema schema, List recordList,
        Connection connection, DwsConfig config) throws DwsClientException, SQLException {
        String sql = JdbcUtil.getBatchSelectFromStatement(schema, schema.getTableName(), recordList);
        BitSet columnMask = new BitSet(schema.getColumns().size());
        for (Get get : recordList) {
            columnMask.or(get.getRecord().getColumnBit());
        }
        ResultSet rs = null;
        try (PreparedStatement ps = connection.prepareStatement(sql)) {
            int paramIndex = 0;
            for (Get get : recordList) {
                Record record = get.getRecord();
                for (int keyIndex : schema.getKeyIndexList()) {
                    ps.setObject(++paramIndex, record.getValue(keyIndex),
                            record.getTableSchema().getColumn(keyIndex).getType());
                }
            }
            ps.setQueryTimeout(
                    config.getReadTimeoutMilliseconds() > 0 ? Math.max(config.getReadTimeoutMilliseconds() / 1000, 1) : 0);
            LogUtil.withLogSwitch(config, () -> log.info("select sql {}", sql));
            rs = ps.executeQuery();
            return parseRecordMap(rs, 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 Map parseRecordMap(ResultSet rs, BitSet columnMask, TableSchema schema)
        throws SQLException {
        Map resultMap = new HashMap<>();
        while (rs != null && rs.next()) {
            Record record = new Record(schema);
            int index = 0;
            for (PrimitiveIterator.OfInt it = columnMask.stream().iterator(); it.hasNext();) {
                int recordColumnIndex = it.next();

                // 将查询结果填充record
                fillRecord(record, recordColumnIndex, rs, ++index, schema.getColumn(recordColumnIndex));
            }
            resultMap.put(record.getRecordKey(), record);
        }
        return resultMap;
    }
}