use of io.stargate.db.dse.impl.StargateClientState in project stargate by stargate.
the class ProxyProtocolQueryInterceptor method interceptQuery.
@Override
public Single<ResultMessage> interceptQuery(CQLStatement statement, QueryState state, QueryOptions options, Map<String, ByteBuffer> customPayload, long queryStartNanoTime) {
// see DseConnection
assert state.getClientState() instanceof StargateClientState;
StargateClientState clientState = (StargateClientState) state.getClientState();
if (!isSystemLocalOrPeers(statement) || !clientState.proxyDestinationAddress().isPresent()) {
return wrapped.map(i -> i.interceptQuery(statement, state, options, customPayload, queryStartNanoTime)).orElse(null);
}
InetSocketAddress destinationAddress = clientState.proxyDestinationAddress().get();
boolean isPrivateDestination = destinationAddress.getAddress().isSiteLocalAddress();
// If the destination is private, we want to use the "source" address of the PROXY header.
// We stored that in clientState.getRemoteAddress().
InetAddress systemLocalAddress = isPrivateDestination ? clientState.getRemoteAddress().getAddress() : destinationAddress.getAddress();
SelectStatement selectStatement = (SelectStatement) statement;
List<List<ByteBuffer>> rows;
String tableName = selectStatement.table();
boolean isPrivateLocal = systemLocalAddress.isSiteLocalAddress();
Set<InetAddress> peers = isPrivateLocal ? privatePeers : publicPeers;
if (tableName.equals(PeersSystemView.NAME)) {
rows = peers.isEmpty() ? Collections.emptyList() : Lists.newArrayListWithCapacity(peers.size() - 1);
for (InetAddress peer : peers) {
if (!peer.equals(systemLocalAddress)) {
rows.add(buildRow(selectStatement.getResultMetadata(), peer, peers));
}
}
} else {
assert tableName.equals(LocalNodeSystemView.NAME);
rows = Collections.singletonList(buildRow(selectStatement.getResultMetadata(), systemLocalAddress, peers));
}
ResultSet resultSet = new ResultSet(selectStatement.getResultMetadata(), rows);
return Single.just(new ResultMessage.Rows(resultSet));
}
use of io.stargate.db.dse.impl.StargateClientState in project stargate by stargate.
the class DefaultQueryInterceptor method interceptSystemLocalOrPeers.
private static Single<ResultMessage> interceptSystemLocalOrPeers(CQLStatement statement, QueryState state, QueryOptions options, long queryStartNanoTime) {
SelectStatement selectStatement = ((SelectStatement) statement);
// Re-parse so that we can intercept and replace the keyspace.
SelectStatement.Raw rawStatement = (SelectStatement.Raw) QueryProcessor.parseStatement(selectStatement.queryString);
rawStatement.setKeyspace(SYSTEM_KEYSPACE_NAME);
SelectStatement interceptStatement = rawStatement.prepare(state.getClientState());
Single<ResultMessage.Rows> rows = interceptStatement.execute(state, options, queryStartNanoTime);
return rows.map(r -> {
// see DseConnection
assert state.getClientState() instanceof StargateClientState;
StargateClientState clientState = (StargateClientState) state.getClientState();
clientState.boundPort().ifPresent(port -> replaceNativeTransportPort(r.result, port));
return new ResultMessage.Rows(new ResultSet(selectStatement.getResultMetadata(), r.result.rows));
});
}
Aggregations