Examples with TokenRangeReadStatementGenerator com.datastax.oss.dsbulk.partitioner.TokenRangeReadStatementGenerator used on opensource projects

Example 1 with TokenRangeReadStatementGenerator

use of com.datastax.oss.dsbulk.partitioner.TokenRangeReadStatementGenerator in project dsbulk by datastax.

the class SchemaSettings method createReadStatements.

public List<Statement<?>> createReadStatements(CqlSession session) {
    PreparedStatement preparedStatement = preparedStatements.get(0);
    ColumnDefinitions variables = preparedStatement.getVariableDefinitions();
    if (variables.size() == 0) {
        return Collections.singletonList(preparedStatement.bind());
    }
    boolean ok = true;
    Optional<CQLWord> start = queryInspector.getTokenRangeRestrictionStartVariable();
    Optional<CQLWord> end = queryInspector.getTokenRangeRestrictionEndVariable();
    if (!start.isPresent() || !end.isPresent()) {
        ok = false;
    }
    if (start.isPresent() && end.isPresent()) {
        Optional<CQLWord> unrecognized = StreamSupport.stream(variables.spliterator(), false).map(columnDefinition -> columnDefinition.getName().asInternal()).map(CQLWord::fromInternal).filter(name -> !name.equals(start.get()) && !name.equals(end.get())).findAny();
        ok = !unrecognized.isPresent();
    }
    if (!ok) {
        throw new IllegalArgumentException("The provided statement (schema.query) contains unrecognized WHERE restrictions; " + "the WHERE clause is only allowed to contain one token range restriction " + "of the form: WHERE token(...) > ? AND token(...) <= ?");
    }
    Metadata metadata = session.getMetadata();
    TokenRangeReadStatementGenerator generator = new TokenRangeReadStatementGenerator(table, metadata);
    List<Statement<?>> statements = generator.generate(splits, range -> preparedStatement.bind().setToken(queryInspector.getTokenRangeRestrictionStartVariableIndex(), range.getStart()).setToken(queryInspector.getTokenRangeRestrictionEndVariableIndex(), range.getEnd()));
    LOGGER.debug("Generated {} bound statements", statements.size());
    // Shuffle the statements to avoid hitting the same replicas sequentially when
    // the statements will be executed.
    Collections.shuffle(statements);
    return statements;
}