com.alibaba.hologres.client.impl.handler.ScanActionHandler Maven / Gradle / Ivy
/*
* Copyright (c) 2022. Alibaba Group Holding Limited
*/
package com.alibaba.hologres.client.impl.handler;
import com.alibaba.hologres.client.EqualsFilter;
import com.alibaba.hologres.client.Filter;
import com.alibaba.hologres.client.HoloConfig;
import com.alibaba.hologres.client.RangeFilter;
import com.alibaba.hologres.client.Scan;
import com.alibaba.hologres.client.exception.HoloClientException;
import com.alibaba.hologres.client.impl.ConnectionHolder;
import com.alibaba.hologres.client.impl.action.ScanAction;
import com.alibaba.hologres.client.model.Record;
import com.alibaba.hologres.client.model.RecordKey;
import com.alibaba.hologres.client.model.RecordScanner;
import com.alibaba.hologres.client.model.TableSchema;
import com.alibaba.hologres.client.utils.IdentifierUtil;
import com.alibaba.hologres.client.utils.Metrics;
import com.codahale.metrics.MetricRegistry;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.Arrays;
import java.util.BitSet;
import java.util.HashMap;
import java.util.Map;
import java.util.PrimitiveIterator;
/**
* ScanAction处理类.
*/
public class ScanActionHandler extends ActionHandler {
public static final Logger LOGGER = LoggerFactory.getLogger(ScanActionHandler.class);
private static final String NAME = "empty";
private final HoloConfig config;
private final ConnectionHolder connectionHolder;
public ScanActionHandler(ConnectionHolder connectionHolder, HoloConfig config) {
super(config);
this.config = config;
this.connectionHolder = connectionHolder;
}
@Override
public void handle(ScanAction scanAction) {
try {
long startTime = System.currentTimeMillis();
Scan scan = scanAction.getScan();
TableSchema schema = scan.getSchema();
BitSet columnMask = new BitSet(schema.getColumnSchema().length);
String tableName = schema.getTableNameObj().getFullName();
if (scan.getSelectedColumns() != null) {
columnMask.or(scan.getSelectedColumns());
} else {
columnMask.set(0, schema.getColumnSchema().length);
}
boolean first = true;
StringBuilder sb = new StringBuilder();
sb.append("select ");
for (PrimitiveIterator.OfInt it = columnMask.stream().iterator(); it.hasNext(); ) {
if (!first) {
sb.append(",");
}
first = false;
sb.append(IdentifierUtil.quoteIdentifier(schema.getColumn(it.next()).getName(), true));
}
sb.append(" from ").append(schema.getTableNameObj().getFullName());
if (scan.getFilterList() != null) {
sb.append(" where ");
first = true;
for (Filter filter : scan.getFilterList()) {
if (!first) {
sb.append(" and ");
}
if (filter instanceof EqualsFilter) {
sb.append(IdentifierUtil.quoteIdentifier(schema.getColumn(((EqualsFilter) filter).getIndex()).getName(), true)).append("=?");
} else if (filter instanceof RangeFilter) {
RangeFilter rf = (RangeFilter) filter;
if (rf.getStart() != null) {
sb.append(IdentifierUtil.quoteIdentifier(schema.getColumn(rf.getIndex()).getName(), true)).append(">=?");
first = false;
}
if (rf.getStop() != null) {
if (!first) {
sb.append(" and ");
}
sb.append(IdentifierUtil.quoteIdentifier(schema.getColumn((rf).getIndex()).getName(), true)).append("");
first = false;
}
}
first = false;
}
}
String[] sortKeyNames = null;
switch (scan.getSortKeys()) {
case PRIMARY_KEY:
if (schema.getPrimaryKeys() != null && schema.getPrimaryKeys().length > 0) {
sortKeyNames = schema.getPrimaryKeys();
}
break;
case CLUSTERING_KEY:
if (schema.getClusteringKey() != null && schema.getClusteringKey().length > 0) {
sortKeyNames = Arrays.stream(schema.getClusteringKey())
// clustering key可能会有":asc"或者":desc"后缀,这里只需要原始字段名
.map(s -> {
if (s.endsWith(":asc")) {
return s.substring(0, s.length() - 4);
} else if (s.endsWith(":desc")) {
return s.substring(0, s.length() - 5);
} else {
return s;
}
})
// 如果字段名被引号引起来就去掉,防止下方拼sql的时候重复加
.map(s -> s.replaceAll("^\"|\"$", ""))
.toArray(String[]::new);
}
break;
case NONE:
default:
}
if (sortKeyNames != null && sortKeyNames.length > 0) {
sb.append(" order by ");
first = true;
for (String name : sortKeyNames) {
if (!first) {
sb.append(",");
}
first = false;
sb.append(IdentifierUtil.quoteIdentifier(name, true));
}
}
String sql = sb.toString();
LOGGER.debug("Scan sql:{}", sql);
connectionHolder.retryExecute((conn) -> {
Map resultMap = new HashMap<>();
try {
if (!config.isUseFixedFe()) {
// AutoCommit set false will execute "begin" query, what fixed fe not support
conn.setAutoCommit(false);
}
try (PreparedStatement ps = conn.prepareStatement(sql, ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY)) {
int paramIndex = 0;
if (scan.getFilterList() != null) {
for (Filter filter : scan.getFilterList()) {
if (filter instanceof EqualsFilter) {
ps.setObject(++paramIndex, ((EqualsFilter) filter).getObj(), schema.getColumn(((EqualsFilter) filter).getIndex()).getType());
} else if (filter instanceof RangeFilter) {
RangeFilter rf = (RangeFilter) filter;
if (rf.getStart() != null) {
ps.setObject(++paramIndex, rf.getStart(), schema.getColumn(rf.getIndex()).getType());
}
if (rf.getStop() != null) {
ps.setObject(++paramIndex, rf.getStop(), schema.getColumn(rf.getIndex()).getType());
}
}
}
}
ps.setFetchSize(scan.getFetchSize() > 0 ? scan.getFetchSize() : config.getScanFetchSize());
ps.setQueryTimeout(scan.getTimeout() > 0 ? scan.getTimeout() : config.getScanTimeoutSeconds());
ResultSet rs = null;
try {
rs = ps.executeQuery();
} catch (SQLException e) {
if (rs != null) {
try {
rs.close();
} catch (SQLException ignore) {
} finally {
rs = null;
}
}
scanAction.getFuture().completeExceptionally(HoloClientException.fromSqlException(e));
}
if (rs != null) {
byte[] lock = new byte[]{};
RecordScanner recordScanner = new RecordScanner(rs, lock, schema, scan.getSelectedColumns());
scanAction.getFuture().complete(recordScanner);
synchronized (lock) {
while (!recordScanner.isDone()) {
try {
lock.wait(5000L);
} catch (InterruptedException e) {
throw new RuntimeException(e);
}
}
}
long endTime = System.currentTimeMillis();
MetricRegistry registry = Metrics.registry();
registry.histogram(Metrics.METRICS_SCAN_LATENCY).update(endTime - startTime);
registry.meter(Metrics.METRICS_SCAN_QPS).mark();
}
}
} catch (SQLException e) {
if (!scanAction.getFuture().isDone()) {
scanAction.getFuture().completeExceptionally(e);
}
throw e;
} finally {
conn.setAutoCommit(true);
}
return null;
}, 1);
} catch (Exception e) {
if (!scanAction.getFuture().isDone()) {
scanAction.getFuture().completeExceptionally(e);
}
}
}
@Override
public String getCostMsMetricName() {
return NAME + METRIC_COST_MS;
}
}